package com.szsh.bigdata.open.utils

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.hbase.client._
import org.apache.hadoop.hbase.filter.{Filter, FilterList, PrefixFilter}
import org.apache.hadoop.hbase.io.ImmutableBytesWritable
import org.apache.hadoop.hbase.mapreduce.{MultiTableOutputFormat, TableInputFormat, TableMapReduceUtil}
import org.apache.hadoop.hbase.util.Bytes
import org.apache.hadoop.hbase.{HBaseConfiguration, TableName}
import org.apache.spark.rdd.PairRDDFunctions
import org.apache.spark.sql.functions.split
import org.apache.spark.sql.{DataFrame, SparkSession}

import java.util

object HbaseProUtils {

    /**
     * 数据写入
     * @param data    待写入数据
     * @param table   数据表
     * @param family  列族
     * @param columns 列名
     * @param host    连接地址
     * @param port    连接端口
     */
    def put(data: DataFrame, table: String, family: String, columns: Array[String], host: String, port: String): Unit = {
        iPut(data, table, family, columns, host, port)
        println(s"写入HBASE: $table 成功")
    }


    /**
     * 删除整行
     * @param data rowkey的df
     * @param table 表名
     * @param host 主机
     * @param port 端口
     */
    def deleteAll(data: DataFrame, table: String, host: String, port: String, cf: String = null, columns: Array[String] = null): Unit = {

        data.repartition(1).foreachPartition(p => {
            val conf = HBaseConfiguration.create()
            conf.set("utils.database.hbase.zookeeper.quorum", host)
            conf.set("utils.database.hbase.zookeeper.property.clientPort", port)
            val connection = ConnectionFactory.createConnection(conf)
            val hTable = connection.getTable(TableName.valueOf(table))
            p.foreach(x => {
                val rowkey = x.getString(0)
                val delete = new Delete(Bytes.toBytes(rowkey))
                hTable.delete(delete)
            })
            hTable.close()
            connection.close()
        })
    }

    /**
     * 删除整列
     * @param table 表名
     * @param host 主机
     * @param port 端口
     * @param columns 列名
     */
    def deleteColumns(table: String, host: String, rowkey: String, cf: String, port: String, columns: Array[String]): Unit = {
        val conf = HBaseConfiguration.create()
        conf.set("utils.database.hbase.zookeeper.quorum", host)
        conf.set("utils.database.hbase.zookeeper.property.clientPort", port)
        val connection = ConnectionFactory.createConnection(conf)
        val hTable = connection.getTable(TableName.valueOf(table))
        val delete = new Delete(rowkey.getBytes())
        import org.apache.hadoop.hbase.util.Bytes
        columns.foreach(x => delete.addColumns(Bytes.toBytes(cf), Bytes.toBytes(x)))
        hTable.delete(delete)
        hTable.close()
        connection.close()
    }


    /**
     * 带前缀过滤的Hbase扫描器
     * @param spark        spark
     * @param table        数据表
     * @param family       列族
     * @param columns      列名
     * @param host         连接地址
     * @param port         连接端口
     * @param prefixFilter 前缀过滤字段
     * @return
     */
    def scanWithPrefixFilter(spark: SparkSession, table: String, family: String,
                             columns: Array[String], host: String, port: String,
                             startRow: String = null, endRow: String = null,
                             prefixFilter: Array[String] = null): DataFrame = {
        iScanWithPrefixFilter(spark, table, family, columns, host, port, startRow, endRow, prefixFilter)
    }


    /**
     * 数据写入，私有
     * @param data          数据
     * @param table         表名
     * @param family        列族
     * @param columns       列名
     * @param host          地址
     * @param port          连接端口
     */
    private[this] def iPut(data: DataFrame, table: String, family: String, columns: Array[String], host: String, port: String): Unit = {
        val conf = HBaseConfiguration.create()
        conf.set("utils.database.hbase.zookeeper.quorum", host)
        conf.set("utils.database.hbase.zookeeper.property.clientPort", port)

        val writeData = data.rdd.repartition(100)
            .map({ row =>
                val rowkey: String = row.getString(0)
                val values: Array[String] = row.getString(1).split("#")
                createHBaseRecord(table, rowkey, family, columns, values)
            }).filter(_ != null)

        new PairRDDFunctions(writeData)
            .saveAsNewAPIHadoopFile(
                "",
                classOf[String],
                classOf[String],
                classOf[MultiTableOutputFormat],
                conf)
    }


    /**
     * 带前缀扫描数据，私有
     * @param spark        spark
     * @param table        数据表
     * @param family       列族
     * @param columns      列名
     * @param host         连接地址
     * @param port         连接端口
     * @param prefixFilter 前缀过滤字段
     * @return
     */
    private[this] def iScanWithPrefixFilter(spark: SparkSession, table: String, family: String, columns: Array[String],
                                            host: String, port: String, startRow: String, endRow: String,
                                            prefixFilter: Array[String]): DataFrame = {
        import spark.implicits._
        val conf: Configuration = HBaseConfiguration.create()
        conf.set("utils.database.hbase.zookeeper.quorum", host)
        conf.set("utils.database.hbase.zookeeper.property.clientPort", port)
        conf.set(TableInputFormat.SCAN_CACHEDROWS, "3000")
        conf.set(TableInputFormat.INPUT_TABLE, table)
        conf.set(TableInputFormat.SCAN_COLUMN_FAMILY, family)
        conf.set(TableInputFormat.SCAN_COLUMNS, columns.toIterable.map(family + ":" + _).mkString(" "))

        val scan = if (startRow == null || endRow == null ) {
            new Scan()
        } else {
            new Scan(startRow.getBytes(), endRow.getBytes())
        }

        if (prefixFilter != null) {
            val filters: util.ArrayList[Filter] = new util.ArrayList()
            prefixFilter.foreach(f => filters.add(new PrefixFilter(f.getBytes())))
            val filterList = new FilterList(FilterList.Operator.MUST_PASS_ONE, filters)
            scan.setFilter(filterList)
            conf.set(TableInputFormat.SCAN, TableMapReduceUtil.convertScanToString(scan))
        }
        
        var data = spark.sparkContext.newAPIHadoopRDD(conf,
            classOf[TableInputFormat],
            classOf[ImmutableBytesWritable],
            classOf[Result]).map(elems => {
            (Bytes.toString(elems._2.getRow), columns.map(col => Bytes.toString(elems._2.getValue(family.getBytes, col.getBytes))).mkString(","))
        }).toDF("rowkey", "infos")
            .withColumn("infos", split($"infos", ","))

        for (i <- columns.indices) {
            data = data.withColumn(s"${columns(i)}", $"infos".getItem(i))
        }
        data
    }


    /**
     * 创建一条记录
     *
     * @param rowkey  -> String类型，为hbase的rowKey
     * @param values  -> String类型，待插入值
     * @param table   -> String类型，为hbase的表名，限制:格式为 "命名空间:表名"
     * @param family  -> String类型，为hbase列族名
     * @param columns -> String类型，为hbase列名，限制:cols的split长度应与values的split长度保持一直
     * @return -> Tuple类型
     */
    private[this] def createHBaseRecord(table: String, rowkey: String, family: String, columns: Array[String], values: Array[String]): (ImmutableBytesWritable, Put) = {
        //列和值应该对应，不对应返回空
        if (columns.length == values.length) {
            // 创建一个put
            val put = new Put(Bytes.toBytes(rowkey))
            columns.iterator.zip(values.iterator).foreach(elems => {
                if (elems._2.nonEmpty) put.addColumn(Bytes.toBytes(family), Bytes.toBytes(elems._1), Bytes.toBytes(elems._2))
            })
            (new ImmutableBytesWritable(Bytes.toBytes(table)), put)
        } else null
    }

}
